# -*- coding: utf-8 -*-

import cv2
import os
import numpy as np
from keras.engine.saving import load_model
from . import faceDataLoad
from photoClassify.project.faceClasiify import cv_puttxt_cn
from photoClassify.project.faceClasiify.config import peoples_dir_path, model_name, video_path, image_size

if __name__ == '__main__':

    # 人脸识别分类器本地存储路径
    cascade_path = "/Users/mc/PycharmProjects/TF-vision/venv/lib/python3.7/site-packages/cv2/data/haarcascade_frontalface_alt2.xml"

    # 使用人脸识别分类器，读入分类器
    cascade = cv2.CascadeClassifier(cascade_path)

    print('---数据加载路径:'+peoples_dir_path)
    print('---模型为:'+model_name)

    nameMap = {}
    id = -1
    dirs = os.listdir(peoples_dir_path)
    dirs.remove(".DS_Store")  # 删除mac文件夹下的隐藏文件
    dirs = sorted(dirs)
    for dir_item in dirs:
        id = id + 1
        nameMap[id] = dir_item

    # 加载模型
    # model.load_model(file_path='./w_g_g1_g2.h5')
    model = load_model('saved_models/' + model_name)

    while True:
        try:
            source = input('输入要测试的视频(文件名+格式)或摄像头(0):')

            # 捕获指定摄像头或视频资源
            if source == '0':
                cap = cv2.VideoCapture(0)
            elif source.__len__() > 0:
                cap = cv2.VideoCapture(video_path + source)
            else:
                continue # 跳过当前循环块中的剩余语句，然后继续进行下一轮循环

            cv2.namedWindow("vision", 0)
            cv2.resizeWindow("vision", 640, 480)
            # 循环检测识别人脸
            while True:
                _, frame = cap.read()  # 读取一帧画面

                # 图像灰化，降低计算复杂度 只是在检测人脸时有效  识别人脸还是用彩色图
                frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

                # 利用分类器检测区域
                faceRects = cascade.detectMultiScale(frame_gray, scaleFactor=1.2, minNeighbors=3, minSize=(40, 40),
                                                     maxSize=(96, 96))
                if len(faceRects) > 0:
                    for faceRect in faceRects:
                        x, y, w, h = faceRect

                        # 识别
                        image = frame[y: y + h, x: x + w]
                        image = cv2.resize(image, image_size)
                        x_valid = np.asarray([np.asarray(image)])
                        x_valid = x_valid.astype('float32') / 255
                        faceID = model.predict(x_valid)
                        if faceID.shape[-1] > 1:
                            result = faceID.argmax(axis=-1)
                        else:
                            result = (faceID > 0.5).astype('int32')
                        cv2.rectangle(frame, (x - 10, y - 10), (x + w + 10, y + h + 10), (101, 101, 255), thickness=2)

                        frame = cv_puttxt_cn.putTxt(frame, nameMap[result[0]], x, y - 40)
                        # cv2.putText(frame, nameMap[result[0]], (x + 30, y + 30), cv2.FONT_HERSHEY_SIMPLEX, 1,
                        #             (101, 101, 255),2)

                cv2.imshow("vision", frame)

                # 快放和慢放
                k = cv2.waitKey(30)
                # 如果输入q则退出循环
                if k & 0xFF == ord('q'):
                    break
            print('视频结束')
            # 释放摄像头并销毁所有窗口
            cap.release()
            cv2.destroyAllWindows()
        except:
            continue
